// handler.cpp
#include <iostream.h>
#include "Component\component.h"
#include "registry.h"

const REG_DATA g_regData[] = {
    { "CLSID\\{11000006-0000-0000-0000-000000000001}", 0, "InsideCOM Handler" },
	{ "CLSID\\{11000006-0000-0000-0000-000000000001}\\InprocHandler32", 0, (const char*)-1 }, 
	{ "CLSID\\{11000006-0000-0000-0000-000000000001}\\InprocHandler32", "ThreadingModel", "Both" },
	{ 0, 0, 0 }
};

const CLSID CLSID_InsideCOMHandler = {0x11000006,0x0000,0x0000,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x01};
HINSTANCE g_hInstance;
long g_cComponents = 0;
long g_cServerLocks = 0;

interface INoAggregationUnknown
{
	virtual HRESULT __stdcall QueryInterface_NoAggregation(REFIID riid, void** ppv)=0;
	virtual ULONG __stdcall AddRef_NoAggregation()=0;
	virtual ULONG __stdcall Release_NoAggregation()=0;
};

class CInsideCOM : public ISum, public INoAggregationUnknown
{
public:
	// IUnknown
	ULONG __stdcall AddRef();
	ULONG __stdcall Release();
	HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

	// INoAggregationUnknown
	ULONG __stdcall AddRef_NoAggregation();
	ULONG __stdcall Release_NoAggregation();
	HRESULT __stdcall QueryInterface_NoAggregation(REFIID riid, void** ppv);

	// ISum
	HRESULT __stdcall Sum(int x, int y, int* retval);

	HRESULT __stdcall Init(IUnknown* pUnknownOuter);

	CInsideCOM() : m_cRef(1) { g_cComponents++; }
	~CInsideCOM()
	{
		cout << "Handler: CInsideCOM::~CInsideCOM()" << endl;
		g_cComponents--;

		// Important: Release the proxy manager in order to release the object
		m_pUnknownInner->Release();
	}

private:
	ULONG m_cRef;
	IUnknown* m_pUnknownInner; // Proxy manager
	IUnknown* m_pUnknownOuter; // Identity object
};

HRESULT CInsideCOM::Init(IUnknown* pUnknownOuter)
{
	// We are aggregated by the identity object (outer)
	m_pUnknownOuter = pUnknownOuter;

	// And we aggregate the proxy manager (inner)
	// For the controlling unknown we simple pass the identity object
	// We could also pass _this_, but that would simply delegate IUnknown calls to the outer object
	return CoGetStdMarshalEx(pUnknownOuter, SMEXF_HANDLER, &m_pUnknownInner);
}

HRESULT CInsideCOM::QueryInterface_NoAggregation(REFIID riid, void** ppv)
{
	if(riid == IID_IUnknown)
		*ppv = (INoAggregationUnknown*)this;
	else if(riid == IID_ISum)
		*ppv = (ISum*)this;

	// IMarshal is provided by the proxy manager (inner)
	else if(riid == IID_IMarshal)
		return m_pUnknownInner->QueryInterface(riid, ppv);

	else 
	{
		*ppv = NULL;
		return E_NOINTERFACE;
	}
	((IUnknown*)(*ppv))->AddRef();
	return S_OK;
}

ULONG CInsideCOM::AddRef_NoAggregation()
{
	return ++m_cRef;
}

ULONG CInsideCOM::Release_NoAggregation()
{
	if(--m_cRef != 0)
		return m_cRef;
	delete this;
	return 0;
}

// These set of IUnknown calls delegate to the identity object
ULONG CInsideCOM::AddRef()
{
	return m_pUnknownOuter->AddRef();
}

ULONG CInsideCOM::Release()
{
	return m_pUnknownOuter->Release();
}

HRESULT CInsideCOM::QueryInterface(REFIID riid, void** ppv)
{
	return m_pUnknownOuter->QueryInterface(riid, ppv);
}

HRESULT CInsideCOM::Sum(int x, int y, int* retval)
{
	HRESULT hr;
	if(x > 50 || y > 50)
	{
		ISum* pSum = 0;
		hr = m_pUnknownInner->QueryInterface(IID_ISum, (void**)&pSum);
		if(FAILED(hr))
			cout << "QueryInterface failed" << endl;
		hr = pSum->Sum(x, y, retval);
		pSum->Release();
		return hr;
	}
	*retval = x + y;
	cout << "Handler: CInsideCOM::Sum() " << x << " + " << y << " = " << *retval << endl;
	return S_OK;
}

class CFactory : public IClassFactory
{
public:
	// IUnknown
	ULONG __stdcall AddRef();
	ULONG __stdcall Release();
	HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

	// IClassFactory
	HRESULT __stdcall CreateInstance(IUnknown *pUnknownOuter, REFIID riid, void** ppv);
	HRESULT __stdcall LockServer(BOOL bLock);

	CFactory() : m_cRef(1) { };
	~CFactory() { };

private:
	long m_cRef;
};

HRESULT CFactory::CreateInstance(IUnknown *pUnknownOuter, REFIID riid, void** ppv)
{
	HRESULT hr;

	if(pUnknownOuter != NULL && riid != IID_IUnknown)
		return CLASS_E_NOAGGREGATION;

	CInsideCOM *pInsideCOM = new CInsideCOM;
	cout << "Handler: CFactory::CreateInstance() " << pInsideCOM << endl;

	if(pInsideCOM == NULL)
		return E_OUTOFMEMORY;

	hr = pInsideCOM->Init(pUnknownOuter);
	if(FAILED(hr))
		cout << "CoGetStdMarshalEx failed" << endl;

	hr = pInsideCOM->QueryInterface_NoAggregation(riid, ppv);
	pInsideCOM->Release_NoAggregation();

	return hr;
}

HRESULT CFactory::LockServer(BOOL bLock)
{
	if(bLock)
		g_cServerLocks++;
	else
		g_cServerLocks--;
	return S_OK;
}

ULONG CFactory::AddRef()
{
	return ++m_cRef;
}

ULONG CFactory::Release()
{
	if(--m_cRef != 0)
		return m_cRef;
	delete this;
	return 0;
}

HRESULT CFactory::QueryInterface(REFIID riid, void** ppv)
{
 	if(riid == IID_IUnknown || riid == IID_IClassFactory)
		*ppv = (IClassFactory*)this;
	else 
	{
		*ppv = NULL;
		return E_NOINTERFACE;
	}
	AddRef();
	return S_OK;
}

HRESULT __stdcall DllCanUnloadNow()
{
	cout << "Handler: DllCanUnloadNow() " << (g_cServerLocks == 0 && g_cComponents == 0 ? "Yes" : "No") << endl;
	if(g_cServerLocks == 0 && g_cComponents == 0)
		return S_OK;
	else
		return S_FALSE;
}

HRESULT __stdcall DllGetClassObject(REFCLSID clsid, REFIID riid, void** ppv)
{
	HRESULT hr;
	
	if(clsid != CLSID_InsideCOMHandler)
		return CLASS_E_CLASSNOTAVAILABLE;

	CFactory* pFactory = new CFactory;
	if(pFactory == NULL)
		return E_OUTOFMEMORY;

	hr = pFactory->QueryInterface(riid, ppv);
	pFactory->Release();
	return hr;
}

HRESULT __stdcall DllRegisterServer()
{
	char DllPath[MAX_PATH];
	GetModuleFileName(g_hInstance, DllPath, sizeof(DllPath));

	return RegisterServerEx(g_regData, DllPath);
}

HRESULT __stdcall DllUnregisterServer()
{
	return UnregisterServerEx(g_regData);
}

BOOL WINAPI DllMain(HINSTANCE hInstance, DWORD dwReason, void* pv)
{
	g_hInstance = hInstance;
	return TRUE;
}